#include "Eigen/Dense"
#include <iostream>
#include <vector>
#include "iomanip"

import numpy as np
from enum import Enum

class MotionState(Enum):
    Stand = 0
    Walk = 1
    Walk2Stand = 2

class LegState(Enum):
    LSt = 0
    RSt = 1
    DSt = 2  # no use but reserved

class DataBus:
    def __init__(self, model_nv):
        self.model_nv = model_nv  # number of dq
        
        # const values for frame mismatch
        self.fe_L_rot_L_off = np.array([[1, 0, 0],
                                        [0, 1, 0],
                                        [0, 0, 1]])
        self.fe_R_rot_L_off = np.array([[1, 0, 0],
                                        [0, 1, 0],
                                        [0, 0, 1]])
        
        # motors, sensors and states feedback
        self.rpy = [0.000] * 3
        self.fL = [0.000] * 3
        self.fR = [0.000] * 3
        self.basePos = [0.000] * 3
        self.baseLinVel = [0.000] * 3  # velocity of the basePos
        self.baseAcc = [0.000] * 3     # baseAcc of the base link
        self.baseAngVel = [0.000] * 3  # angular velocity of the base link
        self.motors_pos_cur = []
        self.motors_vel_cur = []
        self.motors_tor_cur = []
        # self.FL_est = np.array([])
        # self.FR_est = np.array([])
        self.isdqInit = False
        
        # PVT controls
        self.motors_pos_des = []
        self.motors_vel_des = []
        self.motors_tor_des = []
        self.motors_tor_out = []
        
        # states and key variables
        self.q = np.empty(0, dtype=np.float64)
        self.dq = np.empty(0, dtype=np.float64)
        self.ddq = np.empty(0, dtype=np.float64)
        self.qOld = np.empty(0, dtype=np.float64)
        self.J_base = np.empty((0, 0), dtype=np.float64)
        self.J_l = np.empty((0, 0), dtype=np.float64)
        self.J_r = np.empty((0, 0), dtype=np.float64)
        self.J_hd_l = np.empty((0, 0), dtype=np.float64)
        self.J_hd_r = np.empty((0, 0), dtype=np.float64)
        self.J_hip_link = np.empty((0, 0), dtype=np.float64)
        self.dJ_base = np.empty((0, 0), dtype=np.float64)
        self.dJ_l = np.empty((0, 0), dtype=np.float64)
        self.dJ_r = np.empty((0, 0), dtype=np.float64)
        self.dJ_hd_l = np.empty((0, 0), dtype=np.float64)
        self.dJ_hd_r = np.empty((0, 0), dtype=np.float64)
        self.Jcom_W = np.empty((0, 0), dtype=np.float64)  # jacobian of CoM, in world frame
        self.pCoM_W = np.zeros(3)
        self.fe_r_pos_W = np.zeros(3)
        self.fe_l_pos_W = np.zeros(3)
        self.base_pos = np.zeros(3)
        
        self.fe_r_rot_W = np.zeros((3, 3), dtype=np.float64)  # in world frame
        self.fe_l_rot_W = np.zeros((3, 3), dtype=np.float64)
        self.base_rot = np.zeros((3, 3), dtype=np.float64)  # in world frame
        
        self.fe_r_pos_L = np.zeros(3)  # in Body frame
        self.fe_l_pos_L = np.zeros(3)  # in Body frame
        self.hip_link_pos = np.zeros(3)
        self.hip_r_pos_L = np.zeros(3)
        self.hip_l_pos_L = np.zeros(3)
        self.hip_r_pos_W = np.zeros(3)
        self.hip_l_pos_W = np.zeros(3)
        self.fe_r_rot_L = np.zeros((3, 3), dtype=np.float64)
        self.fe_l_rot_L = np.zeros((3, 3), dtype=np.float64)
        self.hip_link_rot = np.zeros((3, 3), dtype=np.float64)
        self.fe_r_pos_L_cmd = np.zeros(3)
        self.fe_l_pos_L_cmd = np.zeros(3)
        self.fe_r_rot_L_cmd = np.zeros((3, 3), dtype=np.float64)
        self.fe_l_rot_L_cmd = np.zeros((3, 3), dtype=np.float64)
        
        self.hd_r_pos_W = np.zeros(3)  # in world frame
        self.hd_l_pos_W = np.zeros(3)  # in world frame
        self.hd_r_rot_W = np.zeros((3, 3), dtype=np.float64)
        self.hd_l_rot_W = np.zeros((3, 3), dtype=np.float64)
        self.hd_r_pos_L = np.zeros(3)  # in body frame
        self.hd_l_pos_L = np.zeros(3)  # in body frame
        self.hd_r_rot_L = np.zeros((3, 3), dtype=np.float64)
        self.hd_l_rot_L = np.zeros((3, 3), dtype=np.float64)
        
        self.qCmd = np.empty(0, dtype=np.float64)
        self.dqCmd = np.empty(0, dtype=np.float64)
        self.tauJointCmd = np.empty(0, dtype=np.float64)
        self.dyn_M = np.empty((0, 0), dtype=np.float64)
        self.dyn_M_inv = np.empty((0, 0), dtype=np.float64)
        self.dyn_C = np.empty((0, 0), dtype=np.float64)
        self.dyn_Ag = np.empty((0, 0), dtype=np.float64)
        self.dyn_dAg = np.empty((0, 0), dtype=np.float64)
        self.dyn_G = np.empty(0, dtype=np.float64)
        self.dyn_Non = np.empty(0, dtype=np.float64)
        self.base_omega_L = np.zeros(3)
        self.base_omega_W = np.zeros(3)
        self.base_rpy = np.zeros(3)
        
        self.slop = np.zeros(3)
        self.inertia = np.zeros((3, 3))
        
        # cmd value from the joystick interpreter
        # self.js_eul_des = np.zeros(3)
        # self.js_pos_des = np.zeros(3)
        # self.js_omega_des = np.zeros(3)
        # self.js_vel_des = np.zeros(3)
        
        # cmd values for MPC
        self.Xd = np.empty(0, dtype=np.float64)
        self.X_cur = np.empty(0, dtype=np.float64)
        # self.mpc_eul_des = np.zeros(3)
        # self.mpc_pos_des = np.zeros(3)
        # self.mpc_omega_des = np.zeros(3)
        # self.mpc_vel_des = np.zeros(3)
        self.X_cal = np.empty(0, dtype=np.float64)
        self.dX_cal = np.empty(0, dtype=np.float64)
        self.fe_react_tau_cmd = np.empty(0, dtype=np.float64)
        
        self.qp_nWSR_MPC = 0
        self.qp_cpuTime_MPC = 0.0
        self.qpStatus_MPC = 0
        
        # cmd values for WBC
        # self.base_rpy_des = np.zeros(3)
        # self.base_pos_des = np.zeros(3)
        self.des_ddq = np.empty(0, dtype=np.float64)
        self.des_dq = np.empty(0, dtype=np.float64)
        self.des_delta_q = np.empty(0, dtype=np.float64)
        self.des_q = np.empty(0, dtype=np.float64)
        self.swing_fe_pos_des_W = np.zeros(3)
        self.swing_fe_rpy_des_W = np.zeros(3)
        self.stance_fe_pos_cur_W = np.zeros(3)
        self.stance_fe_rot_cur_W = np.zeros(3)
        self.wbc_delta_q_final = np.empty(0, dtype=np.float64)
        self.wbc_dq_final = np.empty(0, dtype=np.float64)
        self.wbc_ddq_final = np.empty(0, dtype=np.float64)
        self.wbc_tauJointRes = np.empty(0, dtype=np.float64)
        self.wbc_FrRes = np.empty(0, dtype=np.float64)
        self.Fr_ff = np.empty(0, dtype=np.float64)
        self.qp_nWSR = 0
        self.qp_cpuTime = 0.0
        self.qp_status = 0
        
        # values for foot-placement
        self.swingStartPos_W = np.zeros(3)
        self.swingDesPosCur_W = np.zeros(3)
        self.swingDesPosCur_L = np.zeros(3)
        self.swingDesPosFinal_W = np.zeros(3)
        self.stanceDesPos_W = np.zeros(3)
        self.posHip_W = np.zeros(3)
        self.posST_W = np.zeros(3)
        self.desV_W = np.zeros(3)  # desired linear velocity
        self.desWz_W = 0.0         # desired angular velocity
        self.theta0 = 0.0          # offset yaw angle of the swing leg, w.r.t body frame
        self.width_hips = 0.0      # distance between the left and right hip
        self.tSwing = 0.0
        self.phi = 0.0
        
        # class MotionState(Enum):
        #     Stand = 0
        #     Walk = 1
        #     Walk2Stand = 2
        
        # class LegState(Enum):
        #     LSt = 0
        #     RSt = 1
        #     DSt = 2  # no use but reserved
        
        self.thetaZ_des = 0.0
        self.legState = LegState.LSt
        self.legStateNext = LegState.LSt
        self.motionState = MotionState.Stand
        
        # for jump
        self.base_pos_stand = np.zeros(3)
        self.pfeW_stand = np.zeros((6, 1))
        self.pfeW0 = np.zeros((6, 1))
        
        # Initialize arrays with appropriate sizes
        self.motors_pos_cur = [0] * (self.model_nv - 6)
        self.motors_vel_cur = [0] * (self.model_nv - 6)
        self.motors_tor_out = [0] * (self.model_nv - 6)
        self.motors_tor_cur = [0] * (self.model_nv - 6)
        self.motors_tor_des = [0] * (self.model_nv - 6)
        self.motors_vel_des = [0] * (self.model_nv - 6)
        self.motors_pos_des = [0] * (self.model_nv - 6)
        
        self.q = np.zeros(self.model_nv + 1)
        self.qOld = np.zeros(self.model_nv + 1)
        self.dq = np.zeros(self.model_nv)
        self.ddq = np.zeros(self.model_nv)
        self.qCmd = np.zeros(self.model_nv + 1)
        self.dqCmd = np.zeros(self.model_nv)
        self.tauJointCmd = np.zeros(self.model_nv - 6)
        self.FL_est = np.zeros(6)
        self.FR_est = np.zeros(6)
        self.Xd = np.zeros(12 * 10)
        self.X_cur = np.zeros(12)
        self.X_cal = np.zeros(12)
        self.dX_cal = np.zeros(12)
        self.fe_react_tau_cmd = np.zeros(13 * 3)
        self.Fr_ff = np.zeros(12)
        self.des_ddq = np.zeros(self.model_nv)
        self.des_dq = np.zeros(self.model_nv)
        self.des_delta_q = np.zeros(self.model_nv)
        
        self.base_rpy_des = np.zeros(3)
        self.base_pos_des = np.zeros(3)
        self.js_eul_des = np.zeros(3)
        self.js_pos_des = np.zeros(3)
        self.js_omega_des = np.zeros(3)
        self.js_vel_des = np.zeros(3)
        
        self.motionState = MotionState.Stand

    ## update q according to sensor values, must update sensor values before
    # def updateQ(self, baseAngVel, rpy, motors_pos_cur, motors_vel_cur, baseLinVel, basePos):
    def updateQ(self):
        base_omega_W = np.array([self.baseAngVel[0], self.baseAngVel[1], self.baseAngVel[2]])
        Rcur = self.eul2Rot(self.rpy[0], self.rpy[1], self.rpy[2])
        base_omega_W = Rcur @ base_omega_W
        
        quatNow = self.eul2quat(self.rpy[0], self.rpy[1], self.rpy[2])
        
        self.q[0] = self.basePos[0]
        self.q[1] = self.basePos[1]
        self.q[2] = self.basePos[2]
        self.q[3] = quatNow[0]
        self.q[4] = quatNow[1]
        self.q[5] = quatNow[2]
        self.q[6] = quatNow[3]
        for i in range(len(self.motors_pos_cur)):
            self.q[i + 7] = self.motors_pos_cur[i]
        
        vCoM_W = np.array([self.baseLinVel[0], self.baseLinVel[1], self.baseLinVel[2]])
        self.dq[:3] = vCoM_W
        self.dq[3:6] = base_omega_W
        for i in range(len(self.motors_vel_cur)):
            self.dq[i + 6] = self.motors_vel_cur[i]
        
        self.base_pos = np.array([self.q[0], self.q[1], self.q[2]])
        self.base_rpy = np.array(self.rpy)
        self.base_rot = Rcur
        self.qOld = self.q.copy()
        

    def printdq(q):
        # set 0.00000
        np.set_printoptions(precision=5)
        # print
        print(q[0:6])
        print(q[6:13])
        print(q[13:20])
        print(q[20:24])
        print(q[24:30])
        print(q[30:36])

    def printq(q):
        # set 0.00000
        np.set_printoptions(precision=5)
        # print
        print(q[0:7])
        print(q[7:14])
        print(q[14:21])
        print(q[21:25])
        print(q[25:31])
        print(q[31:37])

    # @staticmethod
    def eul2Rot(self, roll, pitch, yaw):
        # calc
        Rz = np.array([
            [np.cos(yaw), -np.sin(yaw), 0],
            [np.sin(yaw), np.cos(yaw), 0],
            [0, 0, 1]
        ])
        
        Ry = np.array([
            [np.cos(pitch), 0, np.sin(pitch)],
            [0, 1, 0],
            [-np.sin(pitch), 0, np.cos(pitch)]
        ])
        
        Rx = np.array([
            [1, 0, 0],
            [0, np.cos(roll), -np.sin(roll)],
            [0, np.sin(roll), np.cos(roll)]
        ])
        
        return Rz @ Ry @ Rx

    # @staticmethod
    def eul2quat(self, roll, pitch, yaw):
        R = self.eul2Rot(roll, pitch, yaw)
        # Eigen::Quaternion<double> quatCur;
        # quatCur = R; ##rotation matrix converted to quaternion
        # Eigen::Quaterniond resQuat;
        # resQuat=quatCur;
        # return resQuat;
        
        trace = R.trace()
        if trace > 0:
            s = 0.5 / np.sqrt(trace + 1)
            w = 0.25 / s
            x = (R[2, 1] - R[1, 2]) * s
            y = (R[0, 2] - R[2, 0]) * s
            z = (R[1, 0] - R[0, 1]) * s
        else:
            if R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
                s = 2 * np.sqrt(1 + R[0, 0] - R[1, 1] - R[2, 2])
                w = (R[2, 1] - R[1, 2]) / s
                x = 0.25 * s
                y = (R[0, 1] + R[1, 0]) / s
                z = (R[0, 2] + R[2, 0]) / s
            elif R[1, 1] > R[2, 2]:
                s = 2 * np.sqrt(1 + R[1, 1] - R[0, 0] - R[2, 2])
                w = (R[0, 2] - R[2, 0]) / s
                x = (R[0, 1] + R[1, 0]) / s
                y = 0.25 * s
                z = (R[1, 2] + R[2, 1]) / s
            else:
                s = 2 * np.sqrt(1 + R[2, 2] - R[0, 0] - R[1, 1])
                w = (R[1, 0] - R[0, 1]) / s
                x = (R[0, 2] + R[2, 0]) / s
                y = (R[1, 2] + R[2, 1]) / s
                z = 0.25 * s
        
        return np.array([x, y, z, w])
